/*************************************************************************
*																		 *
*				code by Miles Chen (v-mich) 2007.10.8					 *
*																		 *
*																		 *
*************************************************************************/


#pragma once

#include "commdef-int.h"


namespace Noria{

#define MAX_STACK_SIZE			1024
#define COLOR_FRAME				RGB(200,80,0)
#define MAX_SQUARE				1600*1200
#define THREAD_WAIT_TIME		2000
#define INJECT_FUNCTION_SIZE	64

	CSpyCab::CSpyCab(void)
	{
		Hook::EnableDebugPrivilege();

		m_pWnd = (HWND*)Marshal::AllocCoTaskMem(sizeof(HWND)*MAX_STACK_SIZE);
		m_fIsSearch = FALSE;
		m_hPen = CreatePen(PS_SOLID, 3, COLOR_FRAME);
		m_hBrush = (HBRUSH)GetStockObject(HOLLOW_BRUSH);
	}

	CSpyCab::~CSpyCab(void)
	{
		DeleteObject(m_hPen);
		DeleteObject(m_hBrush);
		Marshal::FreeCoTaskMem(m_pWnd);
	}

	int CSpyCab::InitSpy(HWND hDlg, HCURSOR hCursor, 
		int IDC_SPY_HWND, 
		int IDC_SPY_PID, 
		int IDC_SPY_ID, 
		int IDC_SPY_CAPTION, 
		int IDC_SPY_CLASS, 
		int IDC_SPY)
	{
		m_hDlg = hDlg;
		m_hCursorSpy = hCursor;
		this->m_nIDC_SPY_HWND = IDC_SPY_HWND;
		this->m_nIDC_SPY_PID = IDC_SPY_PID;
		this->m_nIDC_SPY_ID = IDC_SPY_ID;
		this->m_nIDC_SPY_CAPTION = IDC_SPY_CAPTION;
		this->m_nIDC_SPY_CLASS = IDC_SPY_CLASS;
		this->m_nIDC_SPY = IDC_SPY;

		return 0;
	}

	int CSpyCab::OnSpyStart()
	{
		//clear
		SetDlgItemTextW(m_hDlg, m_nIDC_SPY_HWND, NULL);
		SetDlgItemTextW(m_hDlg, m_nIDC_SPY_PID, NULL);
		SetDlgItemTextW(m_hDlg, m_nIDC_SPY_ID, NULL);
		SetDlgItemTextW(m_hDlg, m_nIDC_SPY_CAPTION, NULL);
		SetDlgItemTextW(m_hDlg, m_nIDC_SPY_CLASS, NULL);

		m_fIsSearch = TRUE;
		m_hPreWnd = NULL;
		ShowWindow(GetDlgItem(m_hDlg, m_nIDC_SPY), SW_HIDE);
		m_hCursorOriginal = SetCursor(m_hCursorSpy);
		SetCapture(m_hDlg);

		return 0;
	}

	int CSpyCab::OnSpyFinish()
	{
		if(m_fIsSearch)
		{
			SetCursor(m_hCursorOriginal);
			ReleaseCapture();
			RefreshWindow(m_hPreWnd);
			ShowWindow(GetDlgItem(m_hDlg, m_nIDC_SPY), SW_SHOW);
			m_fIsSearch = FALSE;
		}

		return 0;
	}

	int CSpyCab::OnSpySearching()
	{
		if(m_fIsSearch)
		{
			//find window
			GetCursorPos(&m_point);
			m_hWndCurrent = WindowFromPoint(m_point);

			if(!ValidateWindow())
				return 0;

			//draw window
			RefreshWindow(m_hPreWnd);
			HighlightFoundWindow(m_hWndCurrent);

			//HWND
			wsprintfW(m_szBuffer, L"0x%x", m_hWndCurrent);
			SetDlgItemTextW(m_hDlg, m_nIDC_SPY_HWND, m_szBuffer);

			//PID
			if(GetWindowThreadProcessId(m_hWndCurrent, &m_dwProcID))
			{	
				WCHAR szProcBuffer[MAX_CHAR];
				if(Process::GetProcessBaseName(m_dwProcID, szProcBuffer, MAX_CHAR))
					wsprintfW(m_szBuffer, L"%d (%ws)", m_dwProcID, szProcBuffer);
				else
					wsprintfW(m_szBuffer, L"%d (N/A)", m_dwProcID);
			}
			else
			{
				wsprintfW(m_szBuffer, L"(N/A)");
			}
			SetDlgItemTextW(m_hDlg, m_nIDC_SPY_PID, m_szBuffer);

			//automation id
			wsprintfW(m_szBuffer, L"%d", GetDlgCtrlID(m_hWndCurrent));
			SetDlgItemTextW(m_hDlg, m_nIDC_SPY_ID, m_szBuffer);

			//class
			GetClassNameW(m_hWndCurrent, m_szBuffer, MAX_CHAR);
			SetDlgItemTextW(m_hDlg, m_nIDC_SPY_CLASS, m_szBuffer);

			//caption
			if(!SendMessageW(m_hWndCurrent, WM_GETTEXT, MAX_CHAR, (LPARAM)m_szBuffer) && GetLastError() == ERROR_ACCESS_DENIED)
			{
				RemoteGetPasswordEditText(m_hWndCurrent);
			}
			SetDlgItemTextW(m_hDlg, m_nIDC_SPY_CAPTION, m_szBuffer);

			//set new previous window
			m_hPreWnd = m_hWndCurrent;
		}

		return 0;
	}


//private:
	BOOL CSpyCab::ValidateWindow()
	{
		if(!m_hWndCurrent)
			return FALSE;

		if(m_hWndCurrent == m_hDlg || GetParent(m_hWndCurrent) == m_hDlg)
			return FALSE;

		if(!IsWindow(m_hWndCurrent))
			return FALSE;

		//precise find window through point
		PreciseSearch();

		if(m_hWndCurrent == m_hPreWnd)
			return FALSE;

		return TRUE;
	}

	BOOL CSpyCab::EnumChildWindowsEx()
	{
		//init stack
		m_stackPointer = 0;
		m_hWndTemp = NULL;

		while((m_hWndTemp = FindWindowExW(m_hWndSearch, m_hWndTemp, NULL, NULL)) && MAX_STACK_SIZE > m_stackPointer)
		{
			*(m_pWnd + m_stackPointer++) = m_hWndTemp;
		}

		return TRUE;
	}

	int CSpyCab::PreciseSearch()
	{
		if(GetWindow(m_hWndCurrent, GW_CHILD))
		{
			//which means handle is a container, find children
			m_hWndSearch = m_hWndCurrent;		//search children
		}
		else if(GetParent(m_hWndCurrent))
		{
			//if window has a parent
			m_hWndSearch = GetParent(m_hWndCurrent);		//search siblings
		}
		else
			//others... return directly
			return 0;

		//init
		m_hWndBestFound = m_hWndCurrent;
		m_uSmallestWindowSquare = MAX_SQUARE;

		//get children
		EnumChildWindowsEx();

		while(m_stackPointer)
		{
			m_hWndTemp = *(m_pWnd + --m_stackPointer);
			GetWindowRect(m_hWndTemp, &m_rectTemp);
			if(		m_hWndCurrent != m_hWndTemp
				&&	m_rectTemp.left		<= m_point.x
				&&	m_rectTemp.right	>= m_point.x
				&&	m_rectTemp.top		<= m_point.y
				&&	m_rectTemp.bottom	>= m_point.y )
			{
				//compare square, the smaller one is the better one
				if(m_uSmallestWindowSquare > (m_uCurrentWindowSquare = (m_rectTemp.right - m_rectTemp.left)*(m_rectTemp.bottom - m_rectTemp.top)))
				{
					m_hWndBestFound = m_hWndTemp;
					m_uSmallestWindowSquare = m_uCurrentWindowSquare;
				}
			}
		}

		m_hWndCurrent = m_hWndBestFound;

		return 0;
	}

	BOOL CSpyCab::RemoteGetPasswordEditText(HWND hEdit)
	{

		//find process
		if(!GetWindowThreadProcessId(hEdit, &m_dwProcID))
			goto Cleanup;

		//open process
		m_hProc = OpenProcess(PROCESS_CREATE_THREAD | PROCESS_VM_OPERATION | PROCESS_VM_WRITE | PROCESS_VM_READ, TRUE, m_dwProcID);
		if( !m_hProc )
			goto Cleanup;

		// allocate memory
		m_lpRemoteExecute = VirtualAllocEx(m_hProc, NULL, INJECT_FUNCTION_SIZE, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
		if(!m_lpRemoteExecute)
			goto Cleanup;

		m_lpPayload = (LPPAYLOAD)VirtualAllocEx(m_hProc, NULL, sizeof PAYLOAD, MEM_COMMIT, PAGE_READWRITE);
		if(!m_lpPayload)
			goto Cleanup;

		// copy function there
		if(!WriteProcessMemory(m_hProc, m_lpRemoteExecute, &RemoteExecute, INJECT_FUNCTION_SIZE, NULL))
			goto Cleanup;

		// initialize data area for remote guy
		m_payload.hWnd = hEdit;
		m_payload.lpProcSendMessageW = (SendMessageNextW)SendMessageW;
		if(!WriteProcessMemory(m_hProc, m_lpPayload, &m_payload, sizeof PAYLOAD, NULL))
			goto Cleanup;

		// CreateRemoteThread()
		m_hThread = CreateRemoteThread(m_hProc, NULL, 0, (LPTHREAD_START_ROUTINE)m_lpRemoteExecute, m_lpPayload, NULL, NULL);
		if(!m_hThread)
			goto Cleanup;


		m_dwResult = WaitForSingleObject(m_hThread, THREAD_WAIT_TIME);
		switch(m_dwResult)
		{
		case WAIT_TIMEOUT:
			goto Cleanup;
		case WAIT_FAILED:
			goto Cleanup;
		case WAIT_OBJECT_0:
			// this might just have worked, pick up the result!
			if(!ReadProcessMemory(m_hProc, m_lpPayload, &m_payload, sizeof PAYLOAD, NULL))
				goto Cleanup;
			if(m_payload.lpResult)
			{
				String::CopyTo(m_szBuffer, m_payload.lpBuffer);
				String::Append(m_szBuffer, L" (cracked)");
			}
			break;
		default:
			break;
		}


Cleanup:
		if(m_lpRemoteExecute)
			VirtualFreeEx(m_hProc, m_lpRemoteExecute, 0, MEM_RELEASE);
		if(m_lpPayload)
			VirtualFreeEx(m_hProc, m_lpPayload, 0, MEM_RELEASE);
		CloseHandle(m_hProc);

		return TRUE;
	}

	BOOL CSpyCab::HighlightFoundWindow(HWND hWnd)
	{
		// Get the screen coordinates of the rectangle of the found window.
		GetWindowRect(hWnd, &m_rect);

		// Get the window DC of the found window.
		m_hdc = GetWindowDC(hWnd);

		if(m_hdc)
		{
			// Select our created pen into the DC and backup the previous pen.
			m_hPrevPen = (HPEN)SelectObject(m_hdc, m_hPen);

			// Select a transparent brush into the DC and backup the previous brush.
			m_hPrevBrush = (HBRUSH)SelectObject(m_hdc, m_hBrush);

			// Draw a rectangle in the DC covering the entire window area of the found window.
			Rectangle(m_hdc, 0, 0, m_rect.right - m_rect.left, m_rect.bottom - m_rect.top);

			// Reinsert the previous pen and brush into the found window's DC.
			SelectObject(m_hdc, m_hPrevPen);

			SelectObject(m_hdc, m_hPrevBrush);

			// Finally release the DC.
			ReleaseDC(hWnd, m_hdc);
		}

		return TRUE;
	}

	BOOL CSpyCab::RefreshWindow(HWND hWnd)
	{
		HWND hRefreshWnd[] = {hWnd, GetParent(hWnd)};

		for(int i=0; i<2; i++)
		{
			//refresh itself and its parent
			if(hRefreshWnd[i])
			{
				InvalidateRect(hRefreshWnd[i], NULL, TRUE);
				UpdateWindow(hRefreshWnd[i]);
				RedrawWindow(hRefreshWnd[i], NULL, NULL, RDW_FRAME | RDW_INVALIDATE | RDW_UPDATENOW | RDW_ALLCHILDREN);
			}
		}
		return TRUE;
	}

	LRESULT CSpyCab::RemoteExecute(LPPAYLOAD p)
	{
		p->lpResult = p->lpProcSendMessageW(p->hWnd, WM_GETTEXT, MAX_CHAR, (LPARAM)p->lpBuffer);

		return p->lpResult;
	}
}
